Skip to content

Fix sinusoidal_embedding calculation for bf16 precision. #1025

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 1 commit into from
Nov 4, 2025
Merged

Fix sinusoidal_embedding calculation for bf16 precision. #1025

merged 1 commit into from
Nov 4, 2025

Conversation

@krahets
Copy link
Contributor

@krahets krahets commented Oct 31, 2025

The Wan DiT uses bf16 precision during training, so sinusoidal_embedding_1d() is required to return a bf16 tensor. sinusoidal_embedding_1d() casts its result to the dtype of the input timestep. This implies that the input timestep must be a bf16 tensor.

However, timestep is in the range [0, 1000]. Casting it to bf16 causes significant precision loss. For example:

ipdb> self.pipe.scheduler.timesteps[timestep_id]
tensor([998.3897])
ipdb> self.pipe.scheduler.timesteps[timestep_id].to(torch.bfloat16)
tensor([1000.], dtype=torch.bfloat16)

Therefore, timestep should be kept in fp32, and only cast to x.dtype before feeding it into the self.time_embedding layer.

Please let me know if there’s a better approach, thanks!

@gemini-code-assist
Copy link
Contributor

gemini-code-assist bot commented Oct 31, 2025

Summary of Changes

Hello @krahets, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical precision issue encountered during bf16 training of the Wan DiT model. The problem stemmed from premature casting of the timestep value to bf16, leading to numerical inaccuracies in the sinusoidal embeddings. The change ensures that the timestep maintains its fp32 precision throughout the embedding calculation, with the final embedding output being cast to the target bf16 type only when necessary, thereby maintaining data integrity and improving training stability.

Highlights

  • Precision Fix for Sinusoidal Embedding: Addresses a precision loss issue when calculating sinusoidal_embedding_1d for bf16 training in the Wan DiT model. Previously, casting the timestep directly to bf16 caused significant loss of numerical accuracy.
  • Refined Type Casting: The timestep input is now kept in fp32 during the sinusoidal_embedding_1d calculation. The result of this embedding is then explicitly cast to x.dtype (which is bf16 in this context) just before being fed into the self.time_embedding layer, preserving precision while meeting the required input type.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist[bot]
gemini-code-assist bot reviewed Oct 31, 2025
View reviewed changes
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request correctly addresses a precision loss issue when using bf16 precision, as described in the PR description. The change to cast the sinusoidal embedding to x.dtype after its calculation is a clean and effective fix for the problem.

You asked if there was a better approach. While your solution is perfectly fine and localized, a more fundamental fix could be to modify the sinusoidal_embedding_1d function itself to not have its output data type depend on the input's data type (e.g., by always returning a float32 tensor). This would make the function's behavior more predictable. However, since that function is not part of this PR, your change is the most appropriate solution.

Separately, while reviewing the code, I noticed a potential unrelated issue. In WanModel.forward at line 405, self.head(x, t) is called. The t tensor has a shape of (batch_size, dim). Inside Head.forward (line 267), this tensor is added to self.modulation which has a shape of (1, 2, dim). This will likely lead to a broadcasting error at runtime. It seems t might need to be reshaped (e.g., via t.unsqueeze(1)) before being used in the Head module to align the dimensions for addition. This is outside the scope of your current changes but something you may want to investigate in a future PR.

@Artiprocher Artiprocher merged commit 401d7d7 into modelscope:main Nov 4, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@krahets @Artiprocher